# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import logging
import pickle

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from absl import app, flags
import numpy as np

from coin.strategy.marketsim.core import spec_pb2
from coin.strategy.marketsim.tool.result_util import (read_sim_result, collect_fills)
from coin.strategy.marketsim.tool.sim_stat import generate_aggregated_sim_stat
from coin.strategy.marketsim.tool.timed_y.calc import (DEFAULT_TIMED_Y_POINTS,
                                                       generate_timed_y_dataframe)


def get_timed_y(timed_y_df, price_type, return_ratio=True, timed_y_points=None):
  assert price_type in ('mid', 'ask0', 'bid0')
  timed_y_points = timed_y_points or DEFAULT_TIMED_Y_POINTS

  cols = []
  for tp in reversed(timed_y_points):
    cols.append('timed_y/%s.%s' % (price_type, tp))
  cols.append('timed_y/%s' % price_type)
  for tp in timed_y_points:
    cols.append('Y/timed_y/%s.Y%s' % (price_type, tp))

  fill_price = timed_y_df.loc[:, 'timed_y/fill_price']
  fill_sign = timed_y_df.loc[:, 'timed_y/fill_qty'].copy()
  fill_sign.loc[fill_sign > 0.] = 1.
  fill_sign.loc[fill_sign < 0.] = -1.

  timed_y = (timed_y_df.loc[:, cols].subtract(fill_price, axis=0).mul(fill_sign, axis=0))
  if return_ratio:
    timed_y = timed_y.div(fill_price, axis=0)
  timed_y = timed_y.mean(axis=0)
  return list(timed_y)


def setup_plt():
  # plt.style.use('fivethirtyeight')
  # plt.style.use('seaborn-talk')
  # plt.rcParams['figure.dpi'] = 72
  plt.rcParams['savefig.dpi'] = 100
  plt.rcParams['font.family'] = ['Inconsolata']
  plt.rcParams['font.sans-serif'] = ['Inconsolata']
  plt.rcParams['font.monospace'] = ['Inconsolata']
  # plt.rcParams['figure.constrained_layout.use'] = True

  plt.rcParams['lines.linewidth'] = 0.5
  plt.rcParams['xtick.labelsize'] = 7
  plt.rcParams['ytick.labelsize'] = 7


def plot_timed_y(timed_y,
                 time_tick_labels,
                 *,
                 log=False,
                 plot_only=False,
                 ax=None,
                 fig=None,
                 gs=None,
                 fmt='.-',
                 color='black',
                 alpha=1.0,
                 label=None,
                 title=None,
                 legend=False):
  if ax is None:
    gs = gridspec.GridSpec(1, 1, figure=fig) if gs is None else gs.subgridspec(1, 1)
    ax = fig.add_subplot(gs[:, :])

  if log:
    ax.plot(np.log10(np.asarray(timed_y)), fmt, color=color, alpha=alpha, label=label)
  else:
    ax.plot(timed_y, fmt, color=color, alpha=alpha, label=label)

  if plot_only:
    return

  plt.axvline(x=len(time_tick_labels) // 2, color='black', alpha=0.3, linewidth=0.5)
  plt.axhline(y=0, color='black', alpha=0.3, linewidth=0.5)

  ax.set_ylabel('Y-mean')
  ax.set_xticks(list(range(len(time_tick_labels))))
  ax.set_xticklabels(time_tick_labels)
  if title is not None:
    ax.set_title(title)
  if legend:
    ax.legend()

  return ax


def plot_timed_y_dataframe(timed_y_df,
                           fig,
                           draw_bbo=False,
                           gs=None,
                           title=None,
                           timed_y_points=None):
  timed_y_points = timed_y_points or DEFAULT_TIMED_Y_POINTS
  labels = [('-%s' % tp) for tp in reversed(timed_y_points)] + ['0'] + timed_y_points

  gs = gridspec.GridSpec(1, 1, figure=fig) if gs is None else gs.subgridspec(1, 1)
  ax = fig.add_subplot(gs[:, :])

  if draw_bbo:
    ask0_timed_y = get_timed_y(timed_y_df, 'ask0', timed_y_points)
    plot_timed_y(ask0_timed_y,
                 labels,
                 plot_only=True,
                 label='ask0',
                 color='red',
                 fmt=':.',
                 alpha=0.5,
                 ax=ax)

    bid0_timed_y = get_timed_y(timed_y_df, 'bid0', timed_y_points)
    plot_timed_y(bid0_timed_y,
                 labels,
                 plot_only=True,
                 label='bid0',
                 color='blue',
                 fmt=':.',
                 alpha=0.5,
                 ax=ax)

  mid_timed_y = get_timed_y(timed_y_df, 'mid', timed_y_points)
  plot_timed_y(mid_timed_y, labels, label='mid', color='black', ax=ax, title=title, legend=True)


def plot_product_stat_table(product_stat, fig, gs=None, col_width=0.55, highlight_rows=None):
  gs = gs or gridspec.GridSpec(1, 1, figure=fig)
  highlight_rows = (highlight_rows or ('gross_pnl',
                                       'num_orders',
                                       'order_fill_ratio',
                                       'volume',
                                       'fill_ratio',
                                       'taker_ratio',
                                       'return_per_volume'))

  fields = spec_pb2.PerProductSimStatProto.DESCRIPTOR.fields
  rows = [field.name for field in fields]
  row_colors = [('#cccccc' if r in highlight_rows else 'w') for r in rows] * len(rows)
  cell_text = [[getattr(product_stat, field.name)] for field in fields]

  ax = fig.add_subplot(gs)
  ax.axis('off')
  ax.axis('tight')
  table = ax.table(cellText=cell_text,
                   rowLabels=rows,
                   rowColours=row_colors,
                   colLabels=['Value'],
                   colWidths=[col_width],
                   loc='best')
  table.auto_set_font_size(False)
  table.set_fontsize(7.)


def main(args):
  FLAGS = flags.FLAGS
  assert FLAGS.machine

  if FLAGS.pickle_in:
    with open(FLAGS.pickle_in, 'rb') as f:
      df = pickle.load(f)
      sim_stat = pickle.load(f)
      logging.info('Aggregated stat\n%s', sim_stat)

  else:
    sim_results = read_sim_result(args[1:])
    sim_stat = generate_aggregated_sim_stat(sim_results)
    logging.info('Aggregated stat\n%s', sim_stat)

    fills = collect_fills(sim_results=sim_results)

    # Assume only one symbol is traded.
    assert len(fills) == 1
    symbol = list(fills.keys())[0]
    fills = fills[symbol]
    df = generate_timed_y_dataframe(symbol, FLAGS.machine, fills)

    if FLAGS.pickle_out:
      with open(FLAGS.pickle_out, 'wb') as f:
        pickle.dump(df, f)
        pickle.dump(sim_stat, f)

  logging.info('timed_y_dataframe shape: %s', str(df.shape))

  if FLAGS.plot_out:
    buy_mask = df['timed_y/fill_qty'] > 0
    sell_mask = df['timed_y/fill_qty'] < 0

    setup_plt()
    fig = plt.figure(figsize=(10, 10), constrained_layout=True)
    gs = gridspec.GridSpec(3, 2, figure=fig)
    plot_product_stat_table(sim_stat.product_stats[0], fig, gs=gs[:, 0])
    plot_timed_y_dataframe(df, fig, gs=gs[0, 1], title='Timed-Y')
    plot_timed_y_dataframe(df.loc[buy_mask], fig, gs=gs[1, 1], draw_bbo=True, title='Timed-Y Buy')
    plot_timed_y_dataframe(df.loc[sell_mask], fig, gs=gs[2, 1], draw_bbo=True, title='Timed-Y Sell')
    plt.savefig(FLAGS.plot_out)


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')

  flags.DEFINE_string('tag', None, 'Tag regex filter')

  flags.DEFINE_string('machine', None, 'Machine')

  flags.DEFINE_string('pickle_in', None, 'If specified, read from the given pickle file.')

  flags.DEFINE_string('pickle_out', None, 'Picked out')

  flags.DEFINE_string('plot_out', 'timed_y.png', 'Timed Y plot file')
  app.run(main)
